
package net.zschech.gwt.comet.server.impl;

import java.io.IOException;
import java.io.OutputStream;
import java.io.Serializable;
import java.util.Date;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import javax.servlet.ServletContext;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import net.zschech.gwt.comet.server.CometServletResponse;
import net.zschech.gwt.comet.server.CometSession;
import org.atmosphere.cpr.AtmosphereEventLifecycle;
import org.atmosphere.cpr.AtmosphereResource;
import org.atmosphere.cpr.AtmosphereResourceEvent;
import org.atmosphere.cpr.AtmosphereResourceEventListener;
import org.atmosphere.cpr.AtmosphereServlet;
import org.atmosphere.cpr.Serializer;

/**
 *
 * @author p.havelaar
 */
public class AtmosphereAsyncServlet extends AsyncServlet {

    @Override
    public void init(ServletContext context) throws ServletException {
        super.init(context);
    }

    @Override
    public Object suspend(final CometServletResponseImpl response, CometSessionImpl session, HttpServletRequest request) throws IOException {
        
        final AtmosphereResource atm = getAtmosphere(request);

        if (atm != null) {
            resourceMap.put(response, atm);
            if (session != null) {
                session.getHttpSession().setAttribute(AtmosphereServlet.ATMOSPHERE_RESOURCE, atm);
            }
            atm.setSerializer(new Serializer() {
                @Override
                public void write(OutputStream out, Object o) throws IOException {
                    if (o instanceof Serializable) {
                        response.write((Serializable)o);
                    } else  if (o instanceof List) {
                        List<?> list = (List)o;
                        if (list.size() > 0) {
                            if (!(list.get(0) instanceof Serializable)) {
                                throw new IOException("Failed to write a list of objects that are not serializable");
                            }
                            response.write((List<Serializable>)o);
                        }
                    } else {
                        throw new IOException("Failed to write an object that is not serializable");
                    }
                }
            });
            if (atm instanceof AtmosphereEventLifecycle) {
                AtmosphereEventLifecycle ael = (AtmosphereEventLifecycle)atm;
                ael.addEventListener(new AtmosphereResourceEventListener() {

                    @Override
                    public void onResume(AtmosphereResourceEvent<HttpServletRequest, HttpServletResponse> event) {
                        response.tryTerminate();
                    }

                    @Override
                    public void onDisconnect(AtmosphereResourceEvent<HttpServletRequest, HttpServletResponse> event) {
                        try {
                            response.terminate();
                        } catch (IOException e) {
                            log("comet response could not write anymore");
                        }
                    }

                    @Override
                    public void onBroadcast(AtmosphereResourceEvent<HttpServletRequest, HttpServletResponse> event) {
                    }
                });
            }
            atm.suspend();
        } else {
            throw new IllegalStateException("Failed to find Atmosphere resource, unable to suspend.");
        }
        
        return atm;
    }

    @Override
    public void terminate(CometServletResponseImpl response, CometSessionImpl session, boolean serverInitiated, Object suspendInfo) {
        resourceMap.remove(response);
        if (suspendInfo != null) {
            ((AtmosphereResource)suspendInfo).resume();
        }
    }

    @Override
    public void invalidate(CometSessionImpl session) {
        if (session != null) {
            resourceMap.remove(session.getResponse());
            session.getHttpSession().setAttribute(AtmosphereServlet.ATMOSPHERE_RESOURCE, null);
            AtmosphereResource atm = getAtmosphere(session);
            if (atm != null) {
                atm.resume();
            }
        }
    }

    @Override
    public void enqueued(CometSessionImpl session) {
        if (session != null) {
            AtmosphereResource atm = getAtmosphere(session);
            if (checkAlive(atm) && session.isValid()) {
                try {
                    session.getResponse().writeSessionQueue(true);
                } catch (IOException e) {
                    log("Failed to write session queue", e);
                }
            } else {
                session.getResponse().tryTerminate();
            }
        }
    }

	@Override
	public ScheduledFuture<?> scheduleHeartbeat(final CometServletResponseImpl response, CometSessionImpl session) {
		assert Thread.holdsLock(response);
        final AtmosphereResource atm = getAtmosphere(response);
        if (atm == null) {
            response.tryTerminate();
            return null;
        }
        return atm.getBroadcaster().getBroadcasterConfig().getScheduledExecutorService()
                .schedule(new Runnable() {
			@Override
			public void run() {
                if (checkAlive(atm)) {
                    log("Sending heartbeat " + new Date().toString());
                    response.tryHeartbeat();
                } else {
                    response.tryTerminate();
                }
			}
		}, response.getHeartbeat(), TimeUnit.MILLISECONDS);
	}

	@Override
	public ScheduledFuture<?> scheduleSessionKeepAlive(final CometServletResponseImpl response, final CometSessionImpl session) {
		assert Thread.holdsLock(response);
		try {
			HttpSession httpSession = session.getHttpSession();
			if (access(httpSession)) {
				return null;
			}

			long keepAliveTime = response.getSessionKeepAliveScheduleTime();
			if (keepAliveTime <= 0) {
				response.tryTerminate();
				return null;
			}
			else if (keepAliveTime == Long.MAX_VALUE) {
				return null;
			}
			else {
                final AtmosphereResource atm = getAtmosphere(response);
                if (atm == null) {
                    response.tryTerminate();
                    return null;
                }
				return atm.getBroadcaster().getBroadcasterConfig().getScheduledExecutorService()
                        .schedule(new Runnable() {
					@Override
					public void run() {
                        if (checkAlive(atm)) {
                            session.setLastAccessedTime(System.currentTimeMillis());
                            response.scheduleSessionKeepAlive();
                        } else {
                            response.tryTerminate();
                        }
					}
				}, keepAliveTime, TimeUnit.MILLISECONDS);
			}
		}
		catch (IllegalStateException e) {
			// the session has been invalidated
			response.tryTerminate();
			return null;
		}
	}

    public AtmosphereResource getAtmosphere(CometServletResponse response) {
        AtmosphereResource atm = resourceMap.get(response);
        if (atm == null) {
            try {
                atm = getAtmosphere(response.getRequest());
            } catch (IllegalStateException e) {
                return null;
            }
        }
        return atm;
    }
    public AtmosphereResource getAtmosphere(CometSession session) {
        return (AtmosphereResource) session.getHttpSession().getAttribute(AtmosphereServlet.ATMOSPHERE_RESOURCE);
    }

    protected AtmosphereResource getAtmosphere(HttpServletRequest request) {
        return (AtmosphereResource) request.getAttribute(AtmosphereServlet.ATMOSPHERE_RESOURCE);
    }

    protected boolean checkAlive(AtmosphereResource resource) {
        if (resource == null || resource.getBroadcaster() == null) {
            return false;
        }
        Iterator<AtmosphereResource> it = resource.getBroadcaster().getAtmosphereResources();
        while (it.hasNext()) {
            if (it.next() == resource) {
                return true;
            }
        }
        return false;
    }

    private ConcurrentMap<CometServletResponse, AtmosphereResource> resourceMap =
            new ConcurrentHashMap<CometServletResponse, AtmosphereResource>();

}